package com.chengqs.rapid.core.balance;

import com.chengqs.rapid.common.config.ServiceInstance;
import com.chengqs.rapid.common.util.TimeUtil;
import com.chengqs.rapid.core.context.AttributeKey;
import com.chengqs.rapid.core.context.RapidContext;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
 * <B>主类名称：</B>AbstractLoadBalance<BR>
 * <B>概要说明：</B>抽象负载均衡类：主要实现预热的功能<BR>
 *
 * @author: chengqs
 * @create: 2025-03-17 00:58:42
 **/
public abstract class AbstractLoadBalance implements LoadBalance {
    @Override
    public ServiceInstance select(RapidContext context) {
        // MATCH_INSTANCES：服务实例列表现在还没有填充，需要LoadBalancePreFilter的时候进行获取并设置
        Set<ServiceInstance> matchInstance = context.getAttribute(AttributeKey.MATCH_INSTANCES);
        if (matchInstance == null || matchInstance.size() == 0) {
            return null;
        }

        List<ServiceInstance> instances = new ArrayList<>(matchInstance);
        if (instances.size() == 1) {
            return instances.get(0);
        }
        ServiceInstance instance = doSelect(context, instances);
        context.putAttribute(AttributeKey.LOAD_BALANCE_INSTANCE, instance);
        return instance;
    }

    /**
     * 子类实现指定的负载均衡策略选择一个服务
     *
     * @param context
     * @param instances
     * @return
     */
    protected abstract ServiceInstance doSelect(RapidContext context, List<ServiceInstance> instances);

    protected static int getWeight(ServiceInstance instance) {
        int weight = instance.getWeight() == null ? LoadBalance.DEFAULT_WEIGHT : instance.getWeight();
        if (weight > 0) {
            // 服务启动注册的时间
            long timestamp = instance.getRegisterTime();
            if (timestamp > 0L) {
                // 服务启动了多久：
                int upTime = (int) (TimeUtil.currentTimeMillis() - timestamp);
                // 默认预热时间 5min
                int warmup = LoadBalance.DEFAULT_WARMUP;
                if (upTime > 0 && upTime < warmup) {
                    weight = calculateWarmupWeight(weight, upTime, warmup);
                }
            }
        }
        return weight;
    }

    /**
     * 计算服务在预热时间内的新权重
     *
     * @param upTime
     * @param warmup
     * @param weight
     * @return
     */
    private static int calculateWarmupWeight(int upTime, int warmup, int weight) {
        // 计算权重 uptime/warmup * weight
        int ww = (int) ((float) upTime / ((float) warmup / (float) weight));
        return ww < 1 ? 1 : (ww > weight ? weight : ww);
    }
}
